# coding: utf-8
from sqlalchemy.orm import Session

from common import create_session, DbBase


def check_session(func):
    def wrapper(*args, **kwargs):
        is_close = False
        session: Session = kwargs.get('session')
        if not session:
            session: Session = create_session()
            is_close = True
        result = func(*args, **kwargs, session=session)
        if is_close:
            session.close()
        return result

    return wrapper


class BaseModel(DbBase):
    __abstract__ = True

    def json(self):
        return {c.name: getattr(self, c.name) for c in self.__table__.columns}